# 导入必要的库
from segment_anything import sam_model_registry, SamAutomaticMaskGenerator, SamPredictor  # Segment Anything模型的相关模块
import cv2  # OpenCV库，用于图像处理
import matplotlib.pyplot as plt  # Matplotlib库，用于绘制图像
import numpy as np  # NumPy库，用于数值计算

# 定义一个函数来显示分割结果（annotations）
def show_anns(anns):
    """
    显示分割结果。参数:- anns: 分割结果列表（每个元素是一个字典，包含分割信息）
    """
    if len(anns) == 0:  # 如果没有分割结果，则直接返回
        return
    # 按照面积大小对分割结果进行排序，从大到小
    sorted_anns = sorted(anns, key=(lambda x: x['area']), reverse=True)
    ax = plt.gca()  # 获取当前的Axes对象
    ax.set_autoscale_on(False)  # 关闭自动缩放
    polygons = []  # 存储多边形
    color = []  # 存储颜色
    for ann in sorted_anns:
        m = ann['segmentation']  # 获取分割掩码
        img = np.ones((m.shape[0], m.shape[1], 3))  # 创建一个与掩码相同大小的图像
        # 随机生成颜色
        color_mask = np.random.random((1, 3)).tolist()[0]
        # 将随机颜色应用到图像的每个通道
        for i in range(3):
            img[:, :, i] = color_mask[i]
        # 在Axes对象上叠加显示分割掩码和颜色
        ax.imshow(np.dstack((img, m * 0.35)))  # 使用透明度为0.35的掩码叠加在原始图像上

# 读取输入图像并转换颜色空间
image = cv2.imread(r"F:\homework\ai\demo\cat.jpg")  # 使用OpenCV读取图像
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)  # 将BGR格式转换为RGB格式

# 设置SAM模型的检查点路径、设备类型和模型类型
sam_checkpoint = 'F:\\homework\\ai\\Model Checkpoints\\sam_vit_h_4b8939.pth'  # 模型权重文件路径
model_type = "default"  # 使用默认的模型类型
device = "cuda"  # 使用GPU进行加速

# 加载并初始化SAM模型
sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)  # 根据模型类型和检查点加载模型
sam.to(device=device)  # 将模型移动到指定设备（CPU或GPU）

# 创建自动掩码生成器实例
mask_generator = SamAutomaticMaskGenerator(sam)

# 使用自动掩码生成器生成分割掩码
masks = mask_generator.generate(image)

# 创建一个新的图形窗口，并显示原始图像和所有分割结果
plt.figure(figsize=(20, 20))  # 创建一个大尺寸的图形窗口
plt.imshow(image)  # 显示原始图像
show_anns(masks)  # 调用自定义函数显示分割结果
plt.axis('off')  # 关闭坐标轴
plt.show()  # 显示图像